﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace MemoryUtil
{
    [Flags]
    public enum AccessMode { Read = 0x01, Write = 0x02 }

    public class ProcessMemory : IDisposable
    {
        private const int PROCESS_WM_READ = 0x0010;
        private const int PROCESS_VM_WRITE = 0x0020;
        private const int PROCESS_VM_OPERATION = 0x0008;

        private bool _isDisposed;
        private IntPtr _procHandle;

        public ProcessMemory(string processName, AccessMode mode)
        {
            ProcessName = processName;
            AccessMode = mode;
        }

        public bool IsOpen
        {
            get { return _procHandle.ToInt32() != 0; }
        }

        public bool TryOpen()
        {
            CheckThrowObjectDisposed();

            try
            {
                Open();
            }
            catch (Exception) { }

            return IsOpen;
        }

        public void Open()
        {
            CheckThrowObjectDisposed();

            //Find the process by name and open it with the right access mode.
            Process foundProcess;
            if (!TryFindProcess(ProcessName, out foundProcess))
                throw new Exception("Process \""+ ProcessName + "\" could not be found.");

            //Figure out what access to use
            int accessMode = 0;
            if ((AccessMode & AccessMode.Read) == AccessMode.Read)
                accessMode |= PROCESS_WM_READ;
            if ((AccessMode & AccessMode.Write) == AccessMode.Write)
                accessMode |= PROCESS_VM_WRITE | PROCESS_VM_OPERATION;

            //Open the process
            _procHandle = OpenProcess(accessMode, false, foundProcess.Id);
            if (!IsOpen)
                throw new System.ComponentModel.Win32Exception(Marshal.GetLastWin32Error(), "Process \"" + ProcessName + "\" could not be opened with the requested access mode.");
        }

        public string ReadNullTerminatedString(Pointer pointer, int maxLength)
        {
            CheckThrowObjectDisposed();

            var effectiveAddress = GetEffectiveAddress(pointer);
            var data = new byte[2];
            var sb = new StringBuilder();

            for (uint i = 0; i < maxLength * sizeof(char); i += sizeof(char))
            {
                ReadBytesIntoArray(effectiveAddress + i, ref data);
                var c = BitConverter.ToChar(data, 0);
                if (c == '\0') break;
                sb.Append(c);
            }

            return sb.ToString();
        }

        public byte ReadByte(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return Read(pointer, sizeof(byte))[0];
        }

        public sbyte ReadSByte(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            unchecked { return (sbyte)Read(pointer, sizeof(byte))[0]; }
        }

        public char ReadChar(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToChar(Read(pointer, sizeof(char)), 0);
        }

        public ushort ReadUInt16(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToUInt16(Read(pointer, sizeof(ushort)), 0);
        }

        public short ReadInt16(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToInt16(Read(pointer, sizeof(short)), 0);
        }

        public int ReadInt32(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToInt32(Read(pointer, sizeof(int)), 0);
        }

        public uint ReadUInt32(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToUInt32(Read(pointer, sizeof(uint)), 0);
        }

        public long ReadInt64(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToInt64(Read(pointer, sizeof(long)), 0);
        }

        public ulong ReadUInt64(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToUInt64(Read(pointer, sizeof(ulong)), 0);
        }

        public float ReadSingle(Pointer pointer)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToSingle(Read(pointer, sizeof(float)), 0);
        }

        public byte[] Read(Pointer pointer, int numBytes)
        {
            CheckThrowObjectDisposed();

            var ret = new byte[numBytes];

            //Read the desired number of bytes out from the effective address
            ReadBytesIntoArray(GetEffectiveAddress(pointer), ref ret);

            return ret;
        }

        public string ReadNullTerminatedString(uint address, int maxLength)
        {
            CheckThrowObjectDisposed();

            var data = new byte[2];
            var sb = new StringBuilder();

            for (uint i = 0; i < maxLength * sizeof(char); i += sizeof(char))
            {
                ReadBytesIntoArray(address + i, ref data);
                var c = BitConverter.ToChar(data, 0);
                if (c == '\0') break;
                sb.Append(c);
            }

            return sb.ToString();
        }

        public byte ReadByte(uint address)
        {
            CheckThrowObjectDisposed();

            return (byte)Read(address, sizeof(byte))[0];
        }

        public sbyte ReadSByte(uint address)
        {
            CheckThrowObjectDisposed();

            unchecked { return (sbyte)Read(address, sizeof(byte))[0]; }
        }

        public char ReadChar(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToChar(Read(address, sizeof(char)), 0);
        }

        public short ReadInt16(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToInt16(Read(address, sizeof(short)), 0);
        }

        public ushort ReadUInt16(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToUInt16(Read(address, sizeof(ushort)), 0);
        }

        public int ReadInt32(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToInt32(Read(address, sizeof(int)), 0);
        }

        public uint ReadUInt32(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToUInt32(Read(address, sizeof(uint)), 0);
        }

        public long ReadInt64(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToInt64(Read(address, sizeof(long)), 0);
        }

        public ulong ReadUInt64(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToUInt64(Read(address, sizeof(ulong)), 0);
        }

        public float ReadSingle(uint address)
        {
            CheckThrowObjectDisposed();

            return BitConverter.ToSingle(Read(address, sizeof(float)), 0);
        }

        public byte[] Read(uint address, int numBytes)
        {
            CheckThrowObjectDisposed();

            var ret = new byte[numBytes];

            //Read the desired number of bytes out from the effective address
            ReadBytesIntoArray(address, ref ret);

            return ret;
        }

        private void ReadBytesIntoArray(uint address, ref byte[] bytes)
        {
            int bytesRead = 0;
            while (bytesRead != bytes.Length)
            {
                int bytesReadThisIteration = 0;
                if (!ReadProcessMemory(_procHandle.ToInt32(), address + (uint)bytesRead, bytes, bytes.Length - bytesRead, ref bytesReadThisIteration))
                {
                    int errorCode = Marshal.GetLastWin32Error();
                    FixupOpenState();
                    throw new System.ComponentModel.Win32Exception(errorCode, "An error occurred while reading the process's memory.");
                }
                bytesRead += bytesReadThisIteration;
            }
        }

        public void Write(Pointer pointer, byte[] bytes)
        {
            CheckThrowObjectDisposed();

            //Write the desired number of bytes
            WriteBytesOutOfArray(GetEffectiveAddress(pointer), bytes);
        }

        public void Write(uint address, byte[] bytes)
        {
            CheckThrowObjectDisposed();

            //Write the desired number of bytes
            WriteBytesOutOfArray(address, bytes);
        }

        private uint GetEffectiveAddress(Pointer pointer)
        {
            var data = new byte[4];
            var address = pointer.Address;

            //Follow the pointer chain to the end
            foreach (var offset in pointer.PointerOffsets)
            {
                //Read the value from memory and add the offset to it
                ReadBytesIntoArray(address, ref data);
                address = BitConverter.ToUInt32(data, 0) + (uint)offset;
            }

            return address;
        }

        private void WriteBytesOutOfArray(uint address, byte[] bytes)
        {
            int bytesWritten = 0;
            while (bytesWritten != bytes.Length)
            {
                int bytesWrittenThisIteration = 0;
                if (!WriteProcessMemory(_procHandle.ToInt32(), address + (uint)bytesWritten, bytes, bytes.Length - bytesWritten, ref bytesWrittenThisIteration))
                {
                    int errorCode = Marshal.GetLastWin32Error();
                    FixupOpenState();
                    throw new System.ComponentModel.Win32Exception(errorCode, "An error occurred while writing to the process's memory.");
                } 
                bytesWritten += bytesWrittenThisIteration;
            }
        }

        /// <summary>
        /// Checks the existence of the process. If it does not exist, the handle is closed 
        /// and invalidated.
        /// </summary>
        private void FixupOpenState()
        {
            Process process;
            if (IsOpen && !TryFindProcess(ProcessName, out process))
            {
                CloseHandle(_procHandle.ToInt32());
                _procHandle = new IntPtr(0);
            }
        }

        private void CheckThrowObjectDisposed()
        {
            if (_isDisposed)
                throw new ObjectDisposedException(this.ToString());
        }

        private bool TryFindProcess(string processName, out Process foundProcess)
        {
            foundProcess = null;
            var foundProcesses = Process.GetProcessesByName(processName);
            if (foundProcesses.Length == 0)
                return false;
            else
            {
                foundProcess = foundProcesses[0];
                return true;
            }
        }

        [DllImport("kernel32.dll")]
        private static extern IntPtr OpenProcess(int dwDesiredAccess, bool bInheritHandle, int dwProcessId);

        [DllImport("kernel32.dll", SetLastError = true)]
        private static extern bool ReadProcessMemory(int hProcess, uint lpBaseAddress, byte[] lpBuffer, int dwSize, ref int lpNumberOfBytesRead);
        
        [DllImport("kernel32.dll", SetLastError = true)]
        private static extern bool WriteProcessMemory(int hProcess, uint lpBaseAddress, byte[] lpBuffer, int dwSize, ref int lpNumberOfBytesWritten);
        
        [DllImport("kernel32.dll")]
        private static extern bool CloseHandle(int handle);

        private void Dispose(bool isNormalDispose)
        {
            if (!_isDisposed)
            {
                if (isNormalDispose)
                {
                    //Release managed resources safely - nothing to do
                }

                if (IsOpen)
                {
                    //Release unmanaged resources
                    CloseHandle(_procHandle.ToInt32());
                }
            }

            _isDisposed = true;
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        ~ProcessMemory()
        {
            Dispose(false);
        }

        public string ProcessName { get; private set; }
        public AccessMode AccessMode { get; private set; }
    }
}
